# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import math

import torch

from torchvision import transforms
# from torchvision.transforms import functional as F
from torchvision.transforms.v2 import functional as F

import torch
import math
import torchvision.transforms as transforms
import torchvision.transforms.functional as F


class RandomResizedCrop(transforms.RandomResizedCrop):
    """
    用于匹配 TF/TPU 实现的 RandomResizedCrop。
    该实现不使用 for 循环，可能与 torchvision 的实现结果不同。
    参考 BYOL 在 TensorFlow 中的实现：
    https://github.com/deepmind/deepmind-research/blob/master/byol/utils/dataset.py#L206
    """

    @staticmethod
    def get_params(img, scale, ratio):
        """
        计算裁剪的参数，生成随机裁剪的起始位置和大小。

        参数：
            img (PIL.Image 或 Tensor): 输入图像，必须是 PIL 图像或 Tensor 格式。
            scale (tuple): 裁剪区域的面积比例范围，表示裁剪区域相对于图像面积的比例。
            ratio (tuple): 宽高比范围，表示裁剪区域的宽高比。

        返回：
            tuple: 返回裁剪区域的起始位置和大小（i, j, h, w），分别表示裁剪区域的
                   起始行、起始列、高度和宽度。
        """
        # 获取图像的宽度和高度
        width, height = F.get_size(img)  # 通过 torchvision.transforms.functional 获取图像尺寸

        area = height * width  # 计算图像的总面积

        # 根据给定的 scale 区间计算目标裁剪区域的面积
        target_area = area * torch.empty(1).uniform_(scale[0], scale[1]).item()

        # 根据给定的 ratio 区间，计算裁剪区域的宽高比
        log_ratio = torch.log(torch.tensor(ratio))
        aspect_ratio = torch.exp(
            torch.empty(1).uniform_(log_ratio[0], log_ratio[1])
        ).item()

        # 计算裁剪区域的宽度和高度
        w = int(round(math.sqrt(target_area * aspect_ratio)))  # 计算宽度
        h = int(round(math.sqrt(target_area / aspect_ratio)))  # 计算高度

        # 确保裁剪区域的宽度和高度不会超出图像的尺寸
        w = min(w, width)
        h = min(h, height)

        # 随机选择裁剪区域的起始位置 (i, j)
        i = torch.randint(0, height - h + 1, size=(1,)).item()
        j = torch.randint(0, width - w + 1, size=(1,)).item()

        # 返回裁剪区域的起始位置 (i, j) 和裁剪的宽高 (h, w)
        return i, j, h, w
